/*
 *   Copyright 2021 zzh
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 *
 */

package com.aduib.boot.common.compiler.support;

import com.aduib.boot.common.compiler.AbstractCompiler;
import com.aduib.boot.common.compiler.JavaPattern;
import com.aduib.boot.common.exception.CommonException;
import com.aduib.boot.common.util.CharUtils;
import com.aduib.boot.common.util.ClassUtils;
import com.aduib.boot.common.util.IOUtils;
import com.aduib.boot.common.util.StrUtils;
import com.aduib.boot.common.util.URLUtils;

import javax.tools.DiagnosticCollector;
import javax.tools.FileObject;
import javax.tools.ForwardingJavaFileManager;
import javax.tools.JavaCompiler;
import javax.tools.JavaFileManager;
import javax.tools.JavaFileObject;
import javax.tools.JavaFileObject.Kind;
import javax.tools.SimpleJavaFileObject;
import javax.tools.StandardJavaFileManager;
import javax.tools.StandardLocation;
import javax.tools.ToolProvider;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.URI;
import java.net.URL;
import java.net.URLClassLoader;
import java.nio.charset.StandardCharsets;
import java.security.AccessController;
import java.security.PrivilegedAction;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * @description: JdkCompiler
 * @author: zzh
 * @date: 2021/9/14 16:08
 */
public class JdkCompiler extends AbstractCompiler {

  private static final String DEFAULT_JAVA_VERSION = "1.8";
  private final JavaCompiler compiler = ToolProvider.getSystemJavaCompiler();
  private final DiagnosticCollector<JavaFileObject> diagnosticCollector =
      new DiagnosticCollector<>();
  private final AdaptiveJavaFileManager javaFileManager;
  private final AdaptiveClassLoader classLoader;
  private final List<String> options;

  public JdkCompiler(List<String> options) {
    this.options = new ArrayList<>(options);
    StandardJavaFileManager manager =
        compiler.getStandardFileManager(diagnosticCollector, null, null);
    final ClassLoader loader = Thread.currentThread().getContextClassLoader();
    if (loader instanceof URLClassLoader
        && (!"sun.misc.Launcher$AppClassLoader".equals(loader.getClass().getName()))) {
      try {
        URLClassLoader urlClassLoader = (URLClassLoader) loader;
        List<File> files = new ArrayList<File>();
        for (URL url : urlClassLoader.getURLs()) {
          files.add(new File(url.getFile()));
        }
        manager.setLocation(StandardLocation.CLASS_PATH, files);
      } catch (IOException e) {
        throw new CommonException(e.getMessage(), e);
      }
    }
    classLoader =
        AccessController.doPrivileged(
            new PrivilegedAction<AdaptiveClassLoader>() {
              @Override
              public AdaptiveClassLoader run() {
                return new AdaptiveClassLoader(loader);
              }
            });
    javaFileManager = new AdaptiveJavaFileManager(manager, classLoader);
  }

  public JdkCompiler() {
    this(buildDefaultOptions());
  }

  public JdkCompiler(String javaVersion) {
    this(buildDefaultOptions(javaVersion));
  }

  private static List<String> buildDefaultOptions(String javaVersion) {
    return Arrays.asList("-source", javaVersion, "-target", javaVersion);
  }

  private static List<String> buildDefaultOptions() {
    return buildDefaultOptions(DEFAULT_JAVA_VERSION);
  }

  @Override
  public Class<?> doCompile(String name, String sourceCode) throws Throwable {
    int i = name.lastIndexOf('.');
    String packageName = i < 0 ? "" : name.substring(0, i);
    String className = i < 0 ? name : name.substring(i + 1);
    AdaptiveJavaFileObject javaFileObject = new AdaptiveJavaFileObject(className, sourceCode);
    javaFileManager.putFileForInput(
        StandardLocation.SOURCE_PATH,
        packageName,
        className + JavaPattern.JAVA_EXTENSION,
        javaFileObject);
    Boolean result =
        compiler
            .getTask(
                null,
                javaFileManager,
                diagnosticCollector,
                options,
                null,
                Collections.singletonList(javaFileObject))
            .call();
    if (result == null || !result) {
      throw new CommonException(
          "Compilation failed. class: "
              + name
              + ", diagnostics: "
              + diagnosticCollector.toString());
    }
    return classLoader.loadClass(name);
  }

  @Override
  protected Class<?> doCompile(byte[] code, String source) throws Throwable {
    return doCompile(StrUtils.str(code, StandardCharsets.UTF_8), source);
  }

  @Override
  protected Class<?> doCompile(InputStream inputStream, String source) throws Throwable {
    return doCompile(IOUtils.readBytes(inputStream), source);
  }

  private static final class AdaptiveJavaFileObject extends SimpleJavaFileObject {

    private final CharSequence source;
    private ByteArrayOutputStream bytecode;

    public AdaptiveJavaFileObject(final String baseName, final CharSequence source) {
      super(URLUtils.toURI(baseName + JavaPattern.JAVA_EXTENSION), Kind.SOURCE);
      this.source = source;
    }

    AdaptiveJavaFileObject(final String name, final Kind kind) {
      super(URLUtils.toURI(name), kind);
      source = null;
    }

    public AdaptiveJavaFileObject(URI uri, Kind kind) {
      super(uri, kind);
      source = null;
    }

    @Override
    public CharSequence getCharContent(final boolean ignoreEncodingErrors)
        throws UnsupportedOperationException {
      if (source == null) {
        throw new UnsupportedOperationException("source == null");
      }
      return source;
    }

    @Override
    public InputStream openInputStream() {
      return new ByteArrayInputStream(getByteCode());
    }

    @Override
    public OutputStream openOutputStream() {
      return bytecode = new ByteArrayOutputStream();
    }

    public byte[] getByteCode() {
      return bytecode.toByteArray();
    }
  }

  private static final class AdaptiveJavaFileManager
      extends ForwardingJavaFileManager<JavaFileManager> {

    private final AdaptiveClassLoader classLoader;

    private final Map<URI, JavaFileObject> fileObjects = new HashMap<URI, JavaFileObject>();

    public AdaptiveJavaFileManager(JavaFileManager fileManager, AdaptiveClassLoader classLoader) {
      super(fileManager);
      this.classLoader = classLoader;
    }

    @Override
    public FileObject getFileForInput(Location location, String packageName, String relativeName)
        throws IOException {
      FileObject o = fileObjects.get(uri(location, packageName, relativeName));
      if (o != null) {
        return o;
      }
      return super.getFileForInput(location, packageName, relativeName);
    }

    public void putFileForInput(
        StandardLocation location, String packageName, String relativeName, JavaFileObject file) {
      fileObjects.put(uri(location, packageName, relativeName), file);
    }

    private URI uri(Location location, String packageName, String relativeName) {
      return URLUtils.toURI(location.getName() + CharUtils.SLASH + packageName + CharUtils.SLASH + relativeName);
    }

    @Override
    public JavaFileObject getJavaFileForOutput(
        Location location, String qualifiedName, Kind kind, FileObject outputFile)
        throws IOException {
      JavaFileObject file = new AdaptiveJavaFileObject(qualifiedName, kind);
      classLoader.add(qualifiedName, file);
      return file;
    }

    @Override
    public ClassLoader getClassLoader(JavaFileManager.Location location) {
      return classLoader;
    }

    @Override
    public String inferBinaryName(Location loc, JavaFileObject file) {
      if (file instanceof AdaptiveJavaFileObject) {
        return file.getName();
      }
      return super.inferBinaryName(loc, file);
    }

    @Override
    public Iterable<JavaFileObject> list(
        Location location, String packageName, Set<Kind> kinds, boolean recurse)
        throws IOException {
      Iterable<JavaFileObject> result = super.list(location, packageName, kinds, recurse);

      ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader();

      ArrayList<JavaFileObject> files = new ArrayList<JavaFileObject>();

      if (location == StandardLocation.CLASS_PATH && kinds.contains(JavaFileObject.Kind.CLASS)) {
        for (JavaFileObject file : fileObjects.values()) {
          if (file.getKind() == Kind.CLASS && file.getName().startsWith(packageName)) {
            files.add(file);
          }
        }

        files.addAll(classLoader.files());
      } else if (location == StandardLocation.SOURCE_PATH
          && kinds.contains(JavaFileObject.Kind.SOURCE)) {
        for (JavaFileObject file : fileObjects.values()) {
          if (file.getKind() == Kind.SOURCE && file.getName().startsWith(packageName)) {
            files.add(file);
          }
        }
      }

      for (JavaFileObject file : result) {
        files.add(file);
      }

      return files;
    }
  }

  private static final class AdaptiveClassLoader extends ClassLoader {

    private final Map<String, JavaFileObject> classes = new HashMap<String, JavaFileObject>();

    AdaptiveClassLoader(final ClassLoader parentClassLoader) {
      super(parentClassLoader);
    }

    Collection<JavaFileObject> files() {
      return Collections.unmodifiableCollection(classes.values());
    }

    @Override
    protected Class<?> findClass(final String qualifiedClassName) throws ClassNotFoundException {
      JavaFileObject file = classes.get(qualifiedClassName);
      if (file != null) {
        byte[] bytes = ((AdaptiveJavaFileObject) file).getByteCode();
        return defineClass(qualifiedClassName, bytes, 0, bytes.length);
      }
      try {
        return ClassUtils.forName(qualifiedClassName, false, getClass().getClassLoader());
      } catch (ClassNotFoundException nf) {
        return super.findClass(qualifiedClassName);
      }
    }

    void add(final String qualifiedClassName, final JavaFileObject javaFile) {
      classes.put(qualifiedClassName, javaFile);
    }

    @Override
    protected synchronized Class<?> loadClass(final String name, final boolean resolve)
        throws ClassNotFoundException {
      return super.loadClass(name, resolve);
    }

    @Override
    public InputStream getResourceAsStream(final String name) {
      if (name.endsWith(JavaPattern.CLASS_EXTENSION)) {
        String qualifiedClassName =
            name.substring(0, name.length() - JavaPattern.CLASS_EXTENSION.length())
                .replace(File.separatorChar, '.');
        AdaptiveJavaFileObject file = (AdaptiveJavaFileObject) classes.get(qualifiedClassName);
        if (file != null) {
          return IOUtils.toStream(file.getByteCode());
        }
      }
      return super.getResourceAsStream(name);
    }
  }
}
